# /*
#  * @Author: Zhengxuan Ye 
#  * @Date: 2025-07-18 01:57:23 
#  * @Last Modified by:   Zhengxuan Ye 
#  * @Last Modified time: 2025-07-18 01:57:23 
#  */

import sys
import yaml
import os
from PyQt5.QtWidgets import (
    QApplication, QWidget, QLabel, QComboBox, QCheckBox, QPushButton, QVBoxLayout, QHBoxLayout, QGroupBox, QGridLayout, QInputDialog, QMessageBox, QSizePolicy
)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont, QPixmap
from utils import resource_path

CONFIG_PATH = resource_path('config.yaml')
PREVIEW_INPUT = resource_path('images/test.JPG')  # 修正为 images/test.JPG
PREVIEW_OUTPUT = resource_path('preview.jpg')

# 四角水印可选内容
CORNER_OPTIONS = [
    ("相机型号(eg. Nikon Z7)", "Model"),
    ("相机厂商(eg. Nikon)", "Make"),
    ("镜头型号(eg. Nikkor 24-70 f/2.8)", "LensModel"),
    ("拍摄参数(eg. 50mm f/1.8 1/1000s ISO 100)", "Param"),
    ("拍摄时间(eg. 2023-01-01 12:00)", "Datetime"),
    ("拍摄日期(eg. 2023-01-01)", "Date"),
    ("自定义", "Custom"),
    ("无", "None"),
    ("镜头厂商 + 镜头型号", "LensMake_LensModel"),
    ("相机型号 + 镜头型号", "CameraModel_LensModel"),
    ("总像素(MP)", "TotalPixel"),
    ("相机厂商 + 相机型号", "CameraMake_CameraModel"),
    ("文件名", "Filename"),
    ("日期 + 文件名", "Date_Filename"),
    ("日期时间 + 文件名", "Datetime_Filename"),
    ("地理信息", "GeoInfo"),
]

CORNER_LABELS = ["左上角", "左下角", "右上角", "右下角"]
CORNER_KEYS = ["left_top", "left_bottom", "right_top", "right_bottom"]

class WatermarkGUI(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('khalil的水印工具')
        self.resize(900, 700)
        self.setStyleSheet(self.apple_style())
        self.custom_values = ["", "", "", ""]
        self._init_phase = True  # 初始化标志，防止弹窗
        self._last_corner_index = [0, 0, 0, 0]  # 记录每个角落上一次的选项索引
        self.init_ui()
        self._init_phase = False  # 初始化结束

    def apple_style(self):
        return """
        QWidget {
            background: #f7f7fa;
            font-family: 'San Francisco', 'Arial', 'PingFang SC', 'Microsoft YaHei', sans-serif;
            font-size: 18px;
        }
        QGroupBox {
            border: 1.5px solid #e0e0e0;
            border-radius: 16px;
            margin-top: 18px;
            background: #ffffff;
            font-size: 18px;
        }
        QGroupBox:title {
            subcontrol-origin: margin;
            left: 16px;
            padding: 0 6px 0 6px;
            color: #888;
            font-weight: bold;
            font-size: 18px;
        }
        QLabel {
            color: #222;
            font-size: 17px;
        }
        QComboBox {
            border: 1.5px solid #d0d0d0;
            border-radius: 10px;
            padding: 8px 16px;
            background: #fafbfc;
            font-size: 17px;
            min-height: 36px;
        }
        QCheckBox {
            spacing: 12px;
            font-size: 17px;
            min-height: 32px;
        }
        QPushButton {
            background: #007aff;
            color: white;
            border: none;
            border-radius: 10px;
            padding: 12px 32px;
            font-weight: bold;
            font-size: 18px;
            min-width: 120px;
            min-height: 40px;
        }
        QPushButton:hover {
            background: #005ecb;
        }
        QPushButton:pressed {
            background: #003e7e;
        }
        """

    def init_ui(self):
        main_layout = QHBoxLayout()
        main_layout.setSpacing(24)
        main_layout.setContentsMargins(24, 24, 24, 24)

        # 左侧：设置区
        left_layout = QVBoxLayout()
        left_layout.setSpacing(24)

        # 当前设置预览
        self.preview_label = QLabel()
        self.preview_label.setFont(QFont('Arial', 15, QFont.Bold))
        self.preview_label.setStyleSheet('background: #f0f0f3; border-radius: 12px; padding: 18px;')
        self.preview_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
        left_layout.addWidget(self.preview_label)

        # 设置项分组
        settings_group = QGroupBox('设置')
        settings_layout = QGridLayout()
        settings_layout.setHorizontalSpacing(32)
        settings_layout.setVerticalSpacing(18)

        # 1. 布局
        settings_layout.addWidget(QLabel('布局'), 0, 0)
        self.layout_combo = QComboBox()
        self.layout_combo.addItems(['normal(Logo居右)', 'simple(Logo居左)', 'square(Logo居左)', 'right-margin(Logo右侧)'])
        settings_layout.addWidget(self.layout_combo, 0, 1)

        # 2. logo 启用
        self.logo_checkbox = QCheckBox('启用 Logo')
        settings_layout.addWidget(self.logo_checkbox, 1, 0, 1, 2)

        # 3-6. 四角水印（下拉选择）
        self.corner_combos = []
        for i, label in enumerate(CORNER_LABELS):
            settings_layout.addWidget(QLabel(label), 2+i, 0)
            combo = QComboBox()
            for opt in CORNER_OPTIONS:
                combo.addItem(opt[0], opt[1])
            combo.currentIndexChanged.connect(lambda idx, c=i: self.on_corner_changed(c))
            self.corner_combos.append(combo)
            settings_layout.addWidget(combo, 2+i, 1)

        # 7. 默认 logo
        settings_layout.addWidget(QLabel('默认 Logo'), 6, 0)
        self.default_logo_combo = QComboBox()
        self.default_logo_combo.addItems(['FUJIFILM', 'DJI', 'Canon', 'Sony', 'Nikon', 'Leica', 'Olympus', 'Panasonic', 'Pentax', 'Ricoh', 'Hasselblad', 'Apple', 'XMAGE'])
        settings_layout.addWidget(self.default_logo_combo, 6, 1)

        # 8. 更多设置
        more_group = QGroupBox('更多设置')
        more_layout = QVBoxLayout()
        more_layout.setSpacing(10)
        self.more_options = []
        self.more_labels = ['白色边框', '等效焦距', '阴影', '按比例填充']
        for label in self.more_labels:
            cb = QCheckBox(label)
            cb.stateChanged.connect(self.update_preview)
            more_layout.addWidget(cb)
            self.more_options.append(cb)
        more_group.setLayout(more_layout)
        settings_layout.addWidget(more_group, 7, 0, 1, 2)

        settings_group.setLayout(settings_layout)
        left_layout.addWidget(settings_group)

        # 操作按钮
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(32)
        self.save_btn = QPushButton('保存设置并预览')
        self.start_btn = QPushButton('开始处理')
        self.exit_btn = QPushButton('退出')
        self.save_btn.clicked.connect(self.on_save_and_preview)
        self.start_btn.clicked.connect(self.on_start_processing)
        self.exit_btn.clicked.connect(self.close)
        btn_layout.addStretch(1)
        btn_layout.addWidget(self.save_btn)
        btn_layout.addWidget(self.start_btn)
        btn_layout.addWidget(self.exit_btn)
        left_layout.addLayout(btn_layout)
        left_layout.addStretch(1)

        # 右侧：图片预览区
        right_layout = QVBoxLayout()
        right_layout.setSpacing(16)
        right_layout.setContentsMargins(0, 0, 0, 0)
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setStyleSheet('background: #eaeaea; border-radius: 12px;')
        self.image_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        right_layout.addWidget(self.image_label)
        # right_layout.addStretch(1)  # 移除这行，让图片预览区占满整个右侧

        # 合并左右布局，左右权重一致
        main_layout.addLayout(left_layout, 1)
        main_layout.addLayout(right_layout, 1)
        self.setLayout(main_layout)

        # 设置默认值
        self.layout_combo.setCurrentIndex(0)  # normal(Logo居右)
        self.logo_checkbox.setChecked(True)
        self.corner_combos[0].setCurrentIndex(2)  # 镜头型号
        self.corner_combos[1].setCurrentIndex(1)  # 相机厂商
        self.corner_combos[2].setCurrentIndex(3)  # 拍摄参数
        self.corner_combos[3].setCurrentIndex(6)  # 自定义
        # 默认logo为NIKON
        logo_index = self.default_logo_combo.findText('Nikon', Qt.MatchFixedString)
        if logo_index != -1:
            self.default_logo_combo.setCurrentIndex(logo_index)
        self.custom_values[3] = 'photo by khalil'  # 右下角自定义内容

        self.update_preview()
        self.update_image_preview()

    def on_corner_changed(self, idx):
        combo = self.corner_combos[idx]
        value = combo.currentData()
        if getattr(self, '_init_phase', False):
            self._last_corner_index[idx] = combo.currentIndex()
            self.update_preview()
            return
        if value == "Custom":
            text, ok = QInputDialog.getText(self, "自定义内容", f"请输入{CORNER_LABELS[idx]}的自定义内容：", text=self.custom_values[idx])
            if ok:
                self.custom_values[idx] = text
                self._last_corner_index[idx] = combo.currentIndex()
            else:
                # 恢复到上一次的选项
                combo.blockSignals(True)
                combo.setCurrentIndex(self._last_corner_index[idx])
                combo.blockSignals(False)
        else:
            self._last_corner_index[idx] = combo.currentIndex()
        self.update_preview()

    def update_preview(self):
        layout = self.layout_combo.currentText()
        logo = '启用' if self.logo_checkbox.isChecked() else '不启用'
        corners = []
        for i, combo in enumerate(self.corner_combos):
            label = combo.currentText()
            value = combo.currentData()
            if value == "Custom":
                label += f"（{self.custom_values[i]}）"
            corners.append(label)
        default_logo = self.default_logo_combo.currentText()
        more = [label for cb, label in zip(self.more_options, self.more_labels) if cb.isChecked()]
        preview = f"""
当前设置：\n\n布局: {layout}\nlogo: {logo}\n左上角: {corners[0]}\n左下角: {corners[1]}\n右上角: {corners[2]}\n右下角: {corners[3]}\n默认logo: {default_logo}\n更多设置: {', '.join(more) if more else '无'}
"""
        self.preview_label.setText(preview)

    def on_save_and_preview(self):
        # 点击后先显示加载中
        self.image_label.clear()
        self.image_label.setText('正在生成预览...')
        QApplication.processEvents()  # 立即刷新界面
        # 收集所有设置项并写入 config.yaml
        try:
            with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            # 1. 布局类型
            layout_map = {
                'normal(Logo居右)': 'watermark_right_logo',
                'simple(Logo居左)': 'simple',
                'square(Logo居左)': 'square',
                'right-margin(Logo右侧)': 'right-margin',
            }
            config['layout']['type'] = layout_map.get(self.layout_combo.currentText(), 'watermark_right_logo')
            # 2. logo 启用
            config['layout']['logo_enable'] = self.logo_checkbox.isChecked()
            # 3. 四角内容
            for i, key in enumerate(CORNER_KEYS):
                value = self.corner_combos[i].currentData()
                config['layout']['elements'][key]['name'] = value
                if value == 'Custom':
                    config['layout']['elements'][key]['value'] = self.custom_values[i]
                else:
                    config['layout']['elements'][key].pop('value', None)
            # 4. 默认logo
            logo_map = {
                'FUJIFILM': './logos/fujifilm.png',
                'DJI': './logos/DJI.jpg',
                'Canon': './logos/canon.png',
                'Sony': './logos/sony.png',
                'Nikon': './logos/nikon.png',
                'Leica': './logos/leica_logo.png',
                'Olympus': './logos/olympus_blue_gold.png',
                'Panasonic': './logos/panasonic.png',
                'Pentax': './logos/pentax.png',
                'Ricoh': './logos/ricoh.png',
                'Hasselblad': './logos/hasselblad.png',
                'Apple': './logos/apple.png',
                'XMAGE': './logos/xmage.png',
            }
            config['logo']['default']['path'] = logo_map.get(self.default_logo_combo.currentText(), './logos/empty.png')
            # 5. 更多设置
            config['global']['white_margin']['enable'] = self.more_options[0].isChecked()
            config['global']['focal_length']['use_equivalent_focal_length'] = self.more_options[1].isChecked()
            config['global']['shadow']['enable'] = self.more_options[2].isChecked()
            config['global']['padding_with_original_ratio']['enable'] = self.more_options[3].isChecked()
            # 写回 config.yaml
            with open(CONFIG_PATH, 'w', encoding='utf-8') as f:
                yaml.safe_dump(config, f, allow_unicode=True)
        except Exception as e:
            QMessageBox.critical(self, "错误", f"写入配置文件失败：{e}")
            self.update_image_preview()
            return
        # 调用单张图片处理生成预览
        preview_success = False
        try:
            from main import process_single
            process_single(PREVIEW_INPUT, PREVIEW_OUTPUT)
            preview_success = True
        except Exception as e:
            QMessageBox.critical(self, "预览生成失败", f"生成预览图片出错：{e}\n将尝试用原图作为预览。")
            # 尝试用原图作为预览
            try:
                from PIL import Image
                img = Image.open(PREVIEW_INPUT)
                img.save(PREVIEW_OUTPUT)
            except Exception as e2:
                QMessageBox.critical(self, "预览生成失败", f"无法复制原图：{e2}")
        # 无论如何都刷新预览
        self.update_image_preview()

    def generate_preview_image(self):
        # 用当前 config.yaml 设置处理 images/test.JPG，输出 preview.jpg
        try:
            from main import process_single
            process_single(PREVIEW_INPUT, PREVIEW_OUTPUT)
        except ImportError:
            from PIL import Image
            img = Image.open(PREVIEW_INPUT)
            img.save(PREVIEW_OUTPUT)

    def update_image_preview(self):
        # 加载并显示 preview.jpg，等比缩放，不裁剪，完整显示
        if not os.path.exists(PREVIEW_OUTPUT):
            self.image_label.clear()
            self.image_label.setText('暂无预览')
            return
        self.image_label.clear()  # 先清空
        QApplication.processEvents()  # 立即刷新
        # 强制让 QPixmap 重新加载文件，避免缓存
        pixmap = QPixmap()
        pixmap.load(PREVIEW_OUTPUT)
        if pixmap.isNull():
            self.image_label.setText('预览加载失败')
            return
        label_w = self.image_label.width()
        label_h = self.image_label.height()
        pixmap = pixmap.scaled(label_w, label_h, Qt.KeepAspectRatio, Qt.SmoothTransformation)
        self.image_label.setPixmap(pixmap)
        self.image_label.repaint()  # 强制重绘
        QApplication.processEvents()  # 再次刷新

    def resizeEvent(self, event):
        super().resizeEvent(event)
        self.update_image_preview()

    def on_start_processing(self):
        # 右侧显示正在处理
        self.image_label.clear()
        self.image_label.setText('正在批量处理...')
        QApplication.processEvents()
        try:
            from main import process_all_images
            process_all_images()
            QMessageBox.information(self, "处理完成", "所有图片已处理并输出到 output 文件夹！")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"批量处理出错：{e}")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    gui = WatermarkGUI()
    gui.show()
    sys.exit(app.exec_()) 